#!/usr/bin/env python
# -*- encoding: utf-8 -*-

"""
@Author  :   Peike Li
@Contact :   peike.li@yahoo.com
@File    :   datasets.py
@Time    :   8/4/19 3:35 PM
@Desc    :
@License :   This source code is licensed under the license found in the
             LICENSE file in the root directory of this source tree.
"""
import os
import numpy as np
import random
import sys
sys.path.append(os.getcwd())
import random
import torch
import cv2
from torch.utils import data
import albumentations as A
from utils.transforms import get_affine_transform
import torchvision.transforms as transforms
class FaceVideoDateset(data.Dataset):
    def __init__(self, root, dataset, crop_size=[512, 512], transform=None,val=False):
        super(FaceVideoDateset).__init__()
        self.root = root
        self.aspect_ratio = crop_size[1] * 1.0 / crop_size[0]
        self.crop_size = np.asarray(crop_size)
        self.dataset = dataset
        self.scale_factor = 0.5
        self.rotation_factor = 30
        self.face_dir = ['crop175','crop15','crop13','crop1','crop2']
        self.face_seg_dir = ['crop_seg175','crop_seg15','crop_seg13','crop_seg1','crop_seg2']
        self.val = val
        self.flip_prob = 0 #.5
        self.albumentations_tansform = A.Compose([
            A.GaussNoise(p=0.1),
            A.OneOf([
                A.MotionBlur(p=0.2),
                A.MedianBlur(blur_limit=3, p=0.1),
                A.Blur(blur_limit=7, p=0.1),
            ], p=0.2),
            # A.MedianBlur(blur_limit=(30, 60), p=0.05),
            A.HueSaturationValue(p=0.2),
            A.RandomBrightnessContrast(p=0.3),
            # A.RandomBrightnessContrast(brightness_limit=0.5,p=0.05),
            A.RandomContrast(p=0.2),
            A.RandomGamma(p=0.2),
        ])
        if val:
            self.train_path=os.path.join(root, 'val_crop_gaussian')
            self.train_list=[]
        else:
            self.train_path=os.path.join(root, 'train_crop_gaussian')
            self.train_list=[]
        for id in os.listdir(os.path.join(self.train_path, 'crop13')):
            self.train_list += [id+'/'+img for img in os.listdir(os.path.join(self.train_path, 'crop13', id))]
        self.transform = transform
        self.number_samples = len(self.train_list)

    def __len__(self):
        if self.val:
            return self.number_samples
        else:
            return int(self.number_samples/5) #
    def _box2cs(self, box):
        x, y, w, h = box[:4]
        return self._xywh2cs(x, y, w, h)
    def _xywh2cs(self, x, y, w, h):
        center = np.zeros((2), dtype=np.float32)
        center[0] = x + w * 0.5
        center[1] = y + h * 0.5
        if w > self.aspect_ratio * h:
            h = w * 1.0 / self.aspect_ratio
        elif w < self.aspect_ratio * h:
            w = h * self.aspect_ratio
        scale = np.array([w * 1.0, h * 1.0], dtype=np.float32)
        return center, scale

    def _preprocess(self, img, parsing_anno, parsing_anno2):
        #do data augumentation
        sf = self.scale_factor
        rf = self.rotation_factor
        h, w, _ = img.shape
        parsing_anno = parsing_anno
        parsing_anno2=parsing_anno2
        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
        r = 0
        s = s * np.clip(np.random.randn() * sf + 1, 1 - sf, 1 + sf)
        r = np.clip(np.random.randn() * rf, -rf * 2, rf * 2) if random.random() <= 0.6 else 0
        trans = get_affine_transform(person_center, s, r, self.crop_size)
        #flip augmentation
        if random.random() < self.flip_prob:
            img = img[:, ::-1, :]
            parsing_anno = parsing_anno[:, ::-1]
            parsing_anno2 = parsing_anno2[:, ::-1]

            person_center[0] = img.shape[1] - person_center[0] - 1
            (15, 14, 17, 16, 19, 18)
            right_idx = [3,5,12,14]
            left_idx = [2,4,11,13]
            for i in range(0, 4):
                right_pos = np.where(parsing_anno == right_idx[i])
                left_pos = np.where(parsing_anno == left_idx[i])
                parsing_anno[right_pos[0], right_pos[1]] = left_idx[i]
                parsing_anno[left_pos[0], left_pos[1]] = right_idx[i]
                right_pos = np.where(parsing_anno2 == right_idx[i])
                left_pos = np.where(parsing_anno2 == left_idx[i])
                parsing_anno2[right_pos[0], right_pos[1]] = left_idx[i]
                parsing_anno2[left_pos[0], left_pos[1]] = right_idx[i]
        input = cv2.warpAffine(
            img,
            trans,
            (int(self.crop_size[1]), int(self.crop_size[0])),
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(0, 0, 0))
        label_parsing2 = cv2.warpAffine(
            parsing_anno2,
            trans,
            (int(self.crop_size[1]), int(self.crop_size[0])),
            flags=cv2.INTER_NEAREST,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(255))
        label_parsing = cv2.warpAffine(
            parsing_anno,
            trans,
            (int(self.crop_size[1]), int(self.crop_size[0])),
            flags=cv2.INTER_NEAREST,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(255))
        return input, label_parsing,label_parsing2
    def _augmentation(self, im):
        transformed = self.albumentations_tansform(image=im)
        im = transformed["image"]
        return im

    def __getitem__(self, index):
        train_item = self.train_list[index]
        if self.val:
            im_path = os.path.join(self.train_path , 'crop13', train_item)
            parsing_anno_path = os.path.join(self.train_path, 'crop_seg13', train_item).replace('jpg', 'png')
            ann_path = os.path.join(self.train_path, 'crop_seg13', train_item[:-13])#.replace('jpg', 'png')
            imgs = os.listdir(ann_path)
            i_img= random.randint(0,len(imgs)-1)
            parsing_anno_path2=os.path.join(ann_path, imgs[i_img])
        else:
            i_dir = random.randint(0,3)
            im_path = os.path.join(self.train_path , self.face_dir[i_dir], train_item)
            parsing_anno_path = os.path.join(self.train_path, self.face_seg_dir[i_dir], train_item).replace('jpg', 'png')
            ann_path = os.path.join(self.train_path, self.face_seg_dir[i_dir], train_item[:-13])#.replace('jpg', 'png')
            imgs = os.listdir(ann_path)
            i_img= random.randint(0,len(imgs)-1)
            parsing_anno_path2=os.path.join(ann_path, imgs[i_img])
            
        im = cv2.imread(im_path, cv2.IMREAD_COLOR)
        # print(im_path, parsing_anno_path)
        h, w, _ = im.shape
        # parsing_anno = np.zeros((h, w), dtype=np.long)
        parsing_anno = cv2.imread(parsing_anno_path, cv2.IMREAD_GRAYSCALE)
        parsing_anno2=cv2.imread(parsing_anno_path2, cv2.IMREAD_GRAYSCALE)

        #os.listdir(parsing_anno_path[:-14]):
        # meta={}
        meta = {
            'name': im_path,
        }
        im = self._augmentation(im)

        input, parsing_anno, parsing_anno2 =self._preprocess(im, parsing_anno, parsing_anno2)
        parsing_before = np.zeros((18, parsing_anno2.shape[0], parsing_anno2.shape[1]),dtype=np.float32)
        for i in range(18):
            parsing_before[i][parsing_anno2==i]=1
        if self.transform:
            input = self.transform(input)
        # input = 
        # parsing_anno = torch.from_numpy(parsing_anno)
        parsing_before= torch.from_numpy(parsing_before)
        # print(input.shape, parsing_before.shape)
        input = torch.cat((input,parsing_before),dim=0)
        # print(input.shape)
        return input, parsing_anno, meta
def vis_parsing_maps(im, parsing_anno, save_im=False, save_path='',im_name='1.png'):
    stride=0
    # Colors for all 20 parts
    part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
                   [255, 0, 85], [255, 0, 170],
                   [0, 255, 0], [85, 255, 0], [170, 255, 0],
                   [0, 255, 85], [0, 255, 170],
                   [0, 0, 255], [85, 0, 255], [170, 0, 255],
                   [0, 85, 255], [0, 170, 255],
                   [255, 255, 0], [255, 255, 85], [255, 255, 170],
                   [255, 0, 255], [255, 85, 255], [255, 170, 255],
                   [0, 255, 255], [85, 255, 255], [170, 255, 255]]

    im = np.array(im)
    vis_im = im.copy().astype(np.uint8)
    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
    vis_parsing_anno[vis_parsing_anno==255]=0
    # vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255
    num_of_class = np.max(vis_parsing_anno)
    for pi in range(1, num_of_class + 1):
        index = np.where(vis_parsing_anno == pi)
        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]

    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
    vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)

# Save result or not
    if save_im:
        cv2.imwrite(save_path+'/anno/'+im_name[:-4] +'.png', vis_parsing_anno)
        cv2.imwrite(save_path+'/weights_img/'+im_name, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100])

if __name__ == '__main__':
    transform = transforms.Compose([
        transforms.ToTensor(),
        # BGR2RGB_transform(),
        transforms.Normalize(mean=(0,0,0),
                                std=(1,1,1)),
    ])
    dataset = FaceDateset('./data/cvpr', "face_dataset",transform=transform)
    for i in range(len(dataset.train_list)):
        input, parsing_anno, meta= dataset[i]
        print(np.unique(parsing_anno))
        vis_parsing_maps(input, parsing_anno, save_im=True, save_path='./data_process/',im_name=dataset.train_list[i][-12:])